package com.bins.langchain.easy.rag.main;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;

import java.net.URISyntaxException;
import java.net.URL;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.time.Duration;
import java.util.Arrays;
import java.util.List;

public class Rag {

    public static void main(String[] args) {
        URL url = Rag.class.getClassLoader().getResource("笑话.txt");
        Document document = getDocument(url);

        List<Document> documents = Arrays.asList(document);

        InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();
        EmbeddingStoreIngestor.ingest(documents, embeddingStore);

        OllamaChatModel chatModel = OllamaChatModel.builder()
                .baseUrl("http://127.0.0.1:11434")
                .modelName("qwen2")
                .temperature(0d)
                .timeout(Duration.ofSeconds(3000))
                .maxRetries(3)
                .logRequests(true)
                .logResponses(true)
                .build();

        Assistant assistant = AiServices.builder(Assistant.class)
                .chatLanguageModel(chatModel)
                .chatMemory(MessageWindowChatMemory.withMaxMessages(10))
                .contentRetriever(EmbeddingStoreContentRetriever.from(embeddingStore))
                .build();

        String answer = assistant.chat("讲个笑话?");
        System.out.println(answer);

    }

    public static Document getDocument(URL url) {
        Document document = null;
        try {
            Path path = Paths.get(url.toURI());
            document = FileSystemDocumentLoader.loadDocument(path);
        } catch (URISyntaxException e) {

        }
        return document;
    }


}


interface Assistant {
    String chat(String userMessage);
}


